package com.example.ai.uploadImage;

import com.example.demo.ai.entity.ImageGenerationResponse;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.Data;
import okhttp3.*;
import java.io.IOException;
import java.util.*;
import java.util.concurrent.TimeUnit;

public class DouBaoARKClient {
    private static final String ARK_BASE_URL = "https://ark.cn-beijing.volces.com/api/v3/images/generations";
    private static final String ARK_API_KEY = "f2750a40-8441-4aff-8306-af1218808c66"; // 请替换为您的实际API密钥
    private static final MediaType JSON_MEDIA_TYPE = MediaType.parse("application/json; charset=utf-8");

    private final OkHttpClient client;
    private final ObjectMapper objectMapper;

    public DouBaoARKClient() {
        this.client = new OkHttpClient.Builder()
                .connectTimeout(30, TimeUnit.SECONDS) // 连接超时
                .readTimeout(120, TimeUnit.SECONDS)    // 读取超时（图像生成需要较长时间）
                .writeTimeout(30, TimeUnit.SECONDS)    // 写入超时
                .callTimeout(180, TimeUnit.SECONDS)    // 总调用超时
                .build();
        this.objectMapper = new ObjectMapper();
    }

    // 图像生成请求类
    @Data
    static class ImageGenerationRequest {
        private String model = "doubao-seedream-4-0-250828";
        private String prompt;
        private String size = "2K";
        private String sequential_image_generation = "disabled";
        private boolean stream = false;
        private String response_format = "url";
        private boolean watermark = false;
        private String optimize_prompt_options = new HashMap<String,String>().put("mode","fast");

        public ImageGenerationRequest(String prompt) {
            this.prompt = prompt;
        }
    }

    /**
     * 生成图像
     */
    public ImageGenerationResponse generateImage(String prompt) throws IOException {
        ImageGenerationRequest request = new ImageGenerationRequest(prompt);
        return generateImage(request);
    }

    /**
     * 生成图像（使用完整请求对象）
     */
    public ImageGenerationResponse generateImage(ImageGenerationRequest request) throws IOException {
        // 构建请求体
        String requestBodyJson = objectMapper.writeValueAsString(request);

        RequestBody body = RequestBody.create(requestBodyJson, JSON_MEDIA_TYPE);

        Request httpRequest = new Request.Builder()
                .url(ARK_BASE_URL)
                .post(body)
                .addHeader("Authorization", "Bearer " + ARK_API_KEY)
                .addHeader("Content-Type", "application/json")
                .build();

        try (Response response = client.newCall(httpRequest).execute()) {
            if (!response.isSuccessful()) {
                throw new IOException("Unexpected code: " + response.code() + " - " + response.message());
            }

            ResponseBody responseBody = response.body();
            if (responseBody != null) {
                String responseText = responseBody.string();
                return objectMapper.readValue(responseText, ImageGenerationResponse.class);
            } else {
                throw new IOException("Response body is null");
            }
        }
    }
}